{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2f0a7faf6e43d4c1",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Zep Vector Store\n",
    "## A long-term memory store for LLM applications\n",
    "This notebook demonstrates how to use the Zep Vector Store with LlamaIndex.\n",
    "\n",
    "\n",
    "## About Zep\n",
    "Zep makes it easy for developers to add relevant documents, chat history memory & rich user data to their LLM app's prompts. \n",
    "\n",
    "\n",
    "## Note\n",
    "Zep can automatically embed your documents. The LlamaIndex implementation of the Zep Vector Store utilizes LlamaIndex's embedders to do so.\n",
    "\n",
    "\n",
    "## Getting Started\n",
    "**Quick Start Guide:** https://docs.getzep.com/deployment/quickstart/\n",
    "**GitHub:** https://github.com/getzep/zep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bb0627e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install zep-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b6547bd146426a4a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:18.442101Z",
     "start_time": "2023-08-09T18:16:18.426672Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "from uuid import uuid4\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "import os\n",
    "import openai\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e69fa431d972ca88",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:19.876605Z",
     "start_time": "2023-08-09T18:16:18.435711Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
      "NumExpr defaulting to 8 threads.\n"
     ]
    }
   ],
   "source": [
    "from llama_index import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.vector_stores import ZepVectorStore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "43aeb87a167a31e4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:19.883201Z",
     "start_time": "2023-08-09T18:16:19.876847Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# load documents\n",
    "documents = SimpleDirectoryReader(\"../data/paul_graham\").load_data()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5285bd54baf41918",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Create a Zep Vector Store and Index\n",
    "\n",
    "You can use an existing Zep Collection, or create a new one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c8d4cea738a5e867",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:21.181413Z",
     "start_time": "2023-08-09T18:16:19.880363Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: GET http://localhost:8000/healthz \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: GET http://localhost:8000/healthz \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: GET http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c \"HTTP/1.1 404 Not Found\"\n",
      "HTTP Request: GET http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c \"HTTP/1.1 404 Not Found\"\n",
      "INFO:llama_index.vector_stores.zep:Collection graham56f41d07312b46e59df8f68cb5fa820c does not exist, will try creating one with dimensions=1536\n",
      "Collection graham56f41d07312b46e59df8f68cb5fa820c does not exist, will try creating one with dimensions=1536\n",
      "INFO:httpx:HTTP Request: POST http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: GET http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: GET http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c/document \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c/document \"HTTP/1.1 200 OK\"\n"
     ]
    }
   ],
   "source": [
    "from llama_index.storage.storage_context import StorageContext\n",
    "\n",
    "zep_api_url = \"http://localhost:8000\"\n",
    "collection_name = f\"graham{uuid4().hex}\"\n",
    "\n",
    "vector_store = ZepVectorStore(\n",
    "    api_url=zep_api_url, collection_name=collection_name, embedding_dimensions=1536\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c72041211dcbdedc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:24.984643Z",
     "start_time": "2023-08-09T18:16:21.178397Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c/search?limit=2 \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST http://localhost:8000/api/v1/collection/graham56f41d07312b46e59df8f68cb5fa820c/search?limit=2 \"HTTP/1.1 200 OK\"\n",
      "\n",
      "Growing up, the author wrote short stories, experimented with programming on an IBM 1401, nagged his father to buy a TRS-80 computer, wrote simple games, a program to predict how high his model rockets would fly, and a word processor. He also studied philosophy in college, but switched to AI after finding it boring. He was involved in building the infrastructure of the web, and eventually started publishing essays online. He also did some painting and learned to cook for groups.\n"
     ]
    }
   ],
   "source": [
    "query_engine = index.as_query_engine()\n",
    "response = query_engine.query(\"What did the author do growing up?\")\n",
    "\n",
    "print(str(response))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e0732f58bfe12bfd",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Querying with Metadata filters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "356e0ec08af5c035",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:24.997025Z",
     "start_time": "2023-08-09T18:16:24.988234Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from llama_index.schema import TextNode\n",
    "\n",
    "nodes = [\n",
    "    TextNode(\n",
    "        text=\"The Shawshank Redemption\",\n",
    "        metadata={\n",
    "            \"author\": \"Stephen King\",\n",
    "            \"theme\": \"Friendship\",\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"The Godfather\",\n",
    "        metadata={\n",
    "            \"director\": \"Francis Ford Coppola\",\n",
    "            \"theme\": \"Mafia\",\n",
    "        },\n",
    "    ),\n",
    "    TextNode(\n",
    "        text=\"Inception\",\n",
    "        metadata={\n",
    "            \"director\": \"Christopher Nolan\",\n",
    "        },\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "18afce2b836a8dd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:25.267226Z",
     "start_time": "2023-08-09T18:16:24.993331Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: GET http://localhost:8000/healthz \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: GET http://localhost:8000/healthz \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: GET http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625 \"HTTP/1.1 404 Not Found\"\n",
      "HTTP Request: GET http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625 \"HTTP/1.1 404 Not Found\"\n",
      "INFO:llama_index.vector_stores.zep:Collection movies398316dc88a045e7a340c46f6eb52625 does not exist, will try creating one with dimensions=1536\n",
      "Collection movies398316dc88a045e7a340c46f6eb52625 does not exist, will try creating one with dimensions=1536\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625 \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625 \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: GET http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625 \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: GET http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625 \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625/document \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625/document \"HTTP/1.1 200 OK\"\n"
     ]
    }
   ],
   "source": [
    "collection_name = f\"movies{uuid4().hex}\"\n",
    "\n",
    "vector_store = ZepVectorStore(\n",
    "    api_url=zep_api_url, collection_name=collection_name, embedding_dimensions=1536\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex(nodes, storage_context=storage_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "32138298211dddee",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:25.267476Z",
     "start_time": "2023-08-09T18:16:25.264275Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from llama_index.vector_stores.types import ExactMatchFilter, MetadataFilters\n",
    "\n",
    "filters = MetadataFilters(filters=[ExactMatchFilter(key=\"theme\", value=\"Mafia\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4fa27e386c42f215",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:25.900125Z",
     "start_time": "2023-08-09T18:16:25.266520Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625/search?limit=2 \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST http://localhost:8000/api/v1/collection/movies398316dc88a045e7a340c46f6eb52625/search?limit=2 \"HTTP/1.1 200 OK\"\n",
      "\n",
      " The Godfather 0.884073254953458\n"
     ]
    }
   ],
   "source": [
    "retriever = index.as_retriever(filters=filters)\n",
    "result = retriever.retrieve(\"What is inception about?\")\n",
    "\n",
    "for r in result:\n",
    "    print(\"\\n\", r.node.text, r.score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a4b9edf0bf12e984",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-09T18:16:25.907552Z",
     "start_time": "2023-08-09T18:16:25.899867Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
